import time
import os

def run(nproc:int, prog:str, args:str):
    abs_path = os.path.abspath('.')
    abs_prog = os.path.join(abs_path, prog)
    cmd = 'mpirun -np {nproc} -H egpu1:8,egpu2:8,egpu7:8,egpu9:8 {executable_file} {args}'.format(
        nproc = nproc,
        executable_file = abs_prog,
        args = args
    )
    r = os.popen(cmd)
    ret = r.read()
    r.close()
    try:
        cost_time_us = float(ret)
        return cost_time_us
    except:
        print(ret)
        return -1
    assert(0)

    
def benchmark_rand_match():
    sub_len = 10
    min_len = 16
    min_nproc = 1
    max_nproc = 8*4
    nproc = min_nproc
    filename = 'rand_match.csv'
    with open(filename,'w') as f:
        f.write("main_len,sub_len,nproc,time(us)\n")
    while True:
        main_len = min_len
        while True:
            ret = run(nproc, 'rand_match.o', '{} {}'.format(main_len, sub_len))
            if ret == -1 or ret > 1000*1000000:
                break
            res = "{main_len},{sub_len},{nproc},{time}".format(main_len=main_len, sub_len=sub_len, nproc=nproc,time=ret)
            print(res)
            with open(filename,'a') as f:
                f.write(res+'\n')
            main_len*=2
        nproc*=2
        if nproc>max_nproc:
            break

def benchmark_omp_rand_match():
    sub_len = 10
    min_len = 16
    min_nproc = 1
    max_nproc = 8*4
    nproc = min_nproc
    filename = 'omp_rand_match.csv'
    with open(filename,'w') as f:
        f.write("main_len,sub_len,nproc,time(us)\n")
    while True:
        main_len = min_len
        while True:
            ret = run(nproc, 'omp-rand_match.o', '{} {}'.format(main_len, sub_len))
            if ret == -1 or ret > 1000*1000000:
                break
            res = "{main_len},{sub_len},{nproc},{time}".format(main_len=main_len, sub_len=sub_len, nproc=nproc,time=ret)
            print(res)
            with open(filename,'a') as f:
                f.write(res+'\n')
            main_len*=2
        nproc*=2
        if nproc>max_nproc:
            break



def benchmark_cholersky():
    filename = 'cholersky.csv'
    min_len = 2
    min_nproc = 1
    max_nproc = 32
    nproc = min_nproc
    with open(filename,'w') as f:
        f.write('length,nproc,time(us)\n')
    while True:
        length = min_len
        while True:
            ret = run(nproc, 'cholersky.o', '{}'.format(length))
            if ret == -1 or ret > 1000*1000000:
                break
            res = '{},{},{}'.format(length,nproc,ret)
            print(res)
            with open(filename, 'a') as f:
                f.write(res+'\n')
            length*=2
        nproc*=2
        if nproc > max_nproc:
            break

        
    
    

if __name__ == '__main__':
    # ret = run(1,'rand_match.o','1000000 10')
    # benchmark_rand_match()
    # benchmark_cholersky()
    benchmark_omp_rand_match()
